package main

import (
	"ai-gateway/config"
	"ai-gateway/proto"
	"context"
	"fmt"
	"log"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
)

func main() {
	// 定义请求参数
	bizSystem := "test_biz"        // 业务系统标识
	scenario := "text_summary"     // 调用场景
	prompt := "你好，请详细介绍一下自己的功能和特点" // AI请求内容

	// 连接业务服务
	conn, err := grpc.Dial(
		fmt.Sprintf("%s:%d", config.Config.ServiceIP, config.Config.BusinessServicePort),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
	)
	if err != nil {
		log.Fatalf("无法连接到业务服务: %v", err)
	}
	defer conn.Close()

	// 创建客户端
	client := proto.NewBusinessServiceClient(conn)

	// 构造请求
	req := &proto.AIRequest{
		BusinessSystem: bizSystem,
		Scenario:       scenario,
		Prompt:         prompt,
	}
	// 调用流式接口
	stream, err := client.GenerateContentStream(context.Background(), req)
	if err != nil {
		log.Fatalf("调用流式AI服务失败: %v", err)
	}
	// 处理流式响应
	fmt.Printf("AI服务流式响应:\n")
	totalToken := 0
	totalCost := 0
	provider := ""
	for {
		resp, err := stream.Recv()
		if err != nil {
			// 流结束时会收到io.EOF错误
			break
		}
		// 输出当前流片段的结果
		fmt.Printf("%s", resp.Result)
		// 累计统计信息
		totalToken += int(resp.TokenUsed)
		totalCost += int(resp.Cost)
		if provider == "" {
			provider = resp.Provider
		}
	}
	// 输出汇总信息
	fmt.Printf("\n\n使用Token总数: %d\n", totalToken)
	fmt.Printf("总调用成本: %d\n", totalCost)
	fmt.Printf("服务提供商: %s\n", provider)
}
